using Plots


mutable struct Node{T}
    id::Int
    weights::AbstractVector{T}
    bias::T
    batch::AbstractVector{T}
end


Network{T} = AbstractVector{Node{T}}


function get_data(net::Network{T}) where T
    hcat([node.batch for node in net]...)'
end


function create_network(W::AbstractMatrix{T}, b::AbstractVector{T}) where T
    net = Node{eltype(W)}[]
    for i = 1:size(b, 1)
        empty_batch = eltype(W)[]  # initialized to empty.
        push!(net, Node(i, W[:, i], b[i], empty_batch))
    end
    net
end


function reset!(net::Network{T}, data::AbstractMatrix{T}) where T
    for i = 1:size(net, 1)
        net[i].batch = copy(data[i, :])
    end
end


function gibbs_sampling!(net::Network{T}, β::T) where T
    X = get_data(net)  # cache for speeding up.
    Threads.@threads for i = 1:size(net, 1)
        aᵢ = net[i].weights' * X .+ net[i].bias
        net[i].batch = sample_bernoulli.(Flux.σ.(β .* aᵢ)[1,:])
        X[i, :] = net[i].batch  # update cache.
    end
end


"""
Implements the activity rule of Boltzmann machine.
"""
function activate!(net::Network{T}, β::T) where T
    gibbs_sampling!(net, β)
end


"""
Implements the energy per node of the Boltzmann machine state.

Energy is defined as E = -1/2 * ∑ᵢⱼ x[i] * W[i, j] * x[j] - ∑ᵢ b[i] * x[i].
"""
function mean_energy(net::Network{T}) where T
    W = hcat([node.weights for node in net]...)
    b = [node.bias for node in net]
    X = hcat([node.batch for node in net]...)'

    function getdiag(x::AbstractMatrix)
        diag = eltype(x)[]
        for i = 1:size(x, 1)
            push!(diag, x[i, i])
        end
        diag
    end

    (-0.5 .* getdiag(X' * W * X) .- X' * b) ./ size(X, 1)
end


function quantize(x::T, precision::T) where T
    @assert precision > 0
    ratio = x / precision
	quantized = (ratio > 0) ? floor(ratio) : ceil(ratio)
	correction = sign(ratio) * sample_bernoulli(abs(ratio - quantized))
	quantized + correction
end


function quantize(net::Network{T}, sparsity::T) where T<:Real
    W = hcat([node.weights for node in net]...)
    precision = quantile(abs.(collect_nondiag(W)), sparsity)
    _quantize = x -> quantize(x, precision)
    [
        Node(node.id,
             _quantize.(node.weights),
             _quantize(node.bias),
             node.batch)
        for node in net
    ]
end


function display(steps::Int, data::AbstractMatrix{T}, energy::AbstractVector{T}; title=nothing) where T
    data_plot = heatmap(data, title=title, legends=false)
    energy_plot = plot(
        collect(1:size(energy, 1)), energy,
        xlims=(0, steps+1), title="Energy", legends=false)
    plot(
        data_plot, energy_plot;
        layout=(1, 2),
        size=(800, 400),
    )
end
